# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import logging

from coin.base.timestamp import get_timestamp
from coin.exchange.base.tick import FixedTickPrice
from coin.exchange.base.order_gateway import OrderSide, OrderType


class BitmexStopLimitMMExecutor(object):
  def __init__(self, execution_params, order_gateway, logger=None):
    self._logger = logger or logging.getLogger('BitmexStopLimitMMExecutor')
    self._og = order_gateway

    self._product = execution_params['product']
    self._native_symbol = execution_params['product'].native_symbol

    self._lot_size = execution_params['lot_size']
    self._min_pos = execution_params['min_pos']
    self._max_pos = execution_params['max_pos']
    self._post_only = execution_params.get('post_only', True)
    self._order_update_period = (execution_params.get('order_update_period', 2) * (10**9))

    self._tick = FixedTickPrice(0.5)  # TODO(jaewon): Parameterize

    self._pos = 0
    self._sell_on_the_fly_orders = []
    self._sell_untriggered_orders = []
    self._sell_triggered_orders = []
    self._buy_on_the_fly_orders = []
    self._buy_untriggered_orders = []
    self._buy_triggered_orders = []

    self._last_fire_time = 0
    self._create_list = []
    self._amend_list = []

  def _update_orders(self):
    self._pos = self._og.get_position(self._product)

    self._sell_on_the_fly_orders = []
    self._sell_untriggered_orders = []
    self._sell_triggered_orders = []
    self._buy_on_the_fly_orders = []
    self._buy_untriggered_orders = []
    self._buy_triggered_orders = []

    for order in self._og.get_working_order():
      if (order.product != self._native_symbol or not order.order_id.startswith('queue0-')
          or order.internal.trigger_price is None):
        continue

      if order.order_side == OrderSide.SELL:
        if not order.accepted:
          self._sell_on_the_fly_orders.append(order)
        elif not order.internal.triggered:
          self._sell_untriggered_orders.append(order)
        else:
          self._sell_triggered_orders.append(order)

      elif order.order_side == OrderSide.BUY:
        if not order.accepted:
          self._buy_on_the_fly_orders.append(order)
        elif not order.internal.triggered:
          self._buy_untriggered_orders.append(order)
        else:
          self._buy_triggered_orders.append(order)

      else:
        self._logger.error('Unknown side: %s', str(order))
        continue

    sell_key = lambda order: (order.price, order.order_id)
    self._sell_untriggered_orders = sorted(self._sell_untriggered_orders, key=sell_key)
    self._sell_triggered_orders = sorted(self._sell_triggered_orders, key=sell_key)

    buy_key = lambda order: (-order.price, order.order_id)
    self._buy_untriggered_orders = sorted(self._buy_untriggered_orders, key=buy_key)
    self._buy_triggered_orders = sorted(self._buy_triggered_orders, key=buy_key)

  def _cancel_sell_triggered_orders(self, keep_price_min, keep_price_max, max_keep_per_side=1):
    keep_count = 0
    for order in self._sell_triggered_orders:
      if (keep_price_min <= order.price and order.price <= keep_price_max
          and keep_count < max_keep_per_side):
        keep_count += 1
      else:
        self._amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _cancel_buy_triggered_orders(self, keep_price_min, keep_price_max, max_keep_per_side=1):
    keep_count = 0
    for order in self._buy_triggered_orders:
      if (keep_price_min <= order.price and order.price <= keep_price_max
          and keep_count < max_keep_per_side):
        keep_count += 1
      else:
        self._amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _cancel_sell_untriggered_orders(self, exclude_order_id=None):
    for order in self._sell_untriggered_orders:
      if order.order_id == exclude_order_id:
        continue
      self._amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _cancel_buy_untriggered_orders(self, exclude_order_id=None):
    for order in self._buy_untriggered_orders:
      if order.order_id == exclude_order_id:
        continue
      self._amend_list.append(self._og.make_bitmex_amend(order.order_id, leaves_qty=0))

  def _manage_sell_order(self, sell_price, sell_trigger_price):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (sell_price is None or self._sell_on_the_fly_orders or self._pos <= self._min_pos):
      self._cancel_sell_untriggered_orders()
      return

    assert self._pos - self._min_pos > 0, self._pos
    sell_qty = min(self._pos - self._min_pos, self._lot_size)
    assert sell_qty > 0, sell_qty

    # Create new order if there is no working order.
    if not self._sell_untriggered_orders:
      order_id = 'queue0-%d' % get_timestamp()
      order = self._og.make_bitmex_order(symbol=self._native_symbol,
                                         price=sell_price,
                                         qty=sell_qty,
                                         order_side=OrderSide.SELL,
                                         order_type=OrderType.LIMIT,
                                         order_id=order_id,
                                         post_only=self._post_only,
                                         trigger_price=sell_trigger_price,
                                         tag='queue0')
      self._create_list.append(order)
      self._logger.info('SELL-STOPLIMIT %s %f @ %.1f, trigger=%.1f, tag=queue0',
                        self._native_symbol,
                        sell_qty,
                        sell_price,
                        sell_trigger_price)
      return

    # If there is an working order that has the same property, keep it and
    # cancel all other orders.
    for order in self._sell_untriggered_orders:
      if (order.internal.trigger_price == sell_trigger_price
          and order.price == sell_price):  # Relaxed condition. Do not compare qty.
        self._cancel_sell_untriggered_orders(exclude_order_id=order.order_id)
        return

    # Choose one with the smallest trigger price and modify it and
    # cancel all other orders.
    for idx, order in enumerate(self._sell_untriggered_orders):
      if idx == 0:
        chosen_order = order
        continue
      if order.internal.trigger_price < chosen_order.internal.trigger_price:
        chosen_order = order

    self._cancel_sell_untriggered_orders(exclude_order_id=chosen_order.order_id)
    self._amend_list.append(
        self._og.make_bitmex_amend(order_id=chosen_order.order_id,
                                   price=sell_price,
                                   leaves_qty=sell_qty,
                                   trigger_price=sell_trigger_price))
    self._logger.info('SELL-STOPLIMIT %s %f @ %.1f, trigger=%.1f, tag=queue0',
                      self._native_symbol,
                      sell_qty,
                      sell_price,
                      sell_trigger_price)

  def _manage_buy_order(self, buy_price, buy_trigger_price):
    # If there is an on-the-fly order, cancel all working orders and keep the
    # on-the-fly order.
    if (buy_price is None or self._buy_on_the_fly_orders or self._pos >= self._max_pos):
      self._cancel_buy_untriggered_orders()
      return

    assert self._max_pos - self._pos > 0, self._pos
    buy_qty = min(self._max_pos - self._pos, self._lot_size)
    assert buy_qty > 0, buy_qty

    # Create new order if there is no working order.
    if not self._buy_untriggered_orders:
      order_id = 'queue0-%d' % get_timestamp()
      order = self._og.make_bitmex_order(symbol=self._native_symbol,
                                         price=buy_price,
                                         qty=buy_qty,
                                         order_side=OrderSide.BUY,
                                         order_type=OrderType.LIMIT,
                                         order_id=order_id,
                                         post_only=self._post_only,
                                         trigger_price=buy_trigger_price,
                                         tag='queue0')
      self._create_list.append(order)
      self._logger.info('BUY-STOPLIMIT  %s %f @ %.1f, trigger=%.1f, tag=queue0',
                        self._native_symbol,
                        buy_qty,
                        buy_price,
                        buy_trigger_price)
      return

    # If there is an working order that has the same property, keep it and
    # cancel all other orders.
    for order in self._buy_untriggered_orders:
      if (order.internal.trigger_price == buy_trigger_price
          and order.price == buy_price):  # Relaxed condition. Do not compare qty.
        self._cancel_buy_untriggered_orders(exclude_order_id=order.order_id)
        return

    # Choose one with the biggest trigger price and modify it and
    # cancel all other orders.
    for idx, order in enumerate(self._buy_untriggered_orders):
      if idx == 0:
        chosen_order = order
        continue
      if order.internal.trigger_price > chosen_order.internal.trigger_price:
        chosen_order = order

    self._cancel_buy_untriggered_orders(exclude_order_id=chosen_order.order_id)
    self._amend_list.append(
        self._og.make_bitmex_amend(order_id=chosen_order.order_id,
                                   price=buy_price,
                                   leaves_qty=buy_qty,
                                   trigger_price=buy_trigger_price))
    self._logger.info('BUY-STOPLIMIT  %s %f @ %.1f, trigger=%.1f, tag=queue0',
                      self._native_symbol,
                      buy_qty,
                      buy_price,
                      buy_trigger_price)

  def execute(self, create_list, amend_list):
    if get_timestamp() - self._last_fire_time < self._order_update_period:
      return

    if amend_list:
      self._og.amend_bulk(amend_list)
      self._logger.debug('Amend %d orders', len(amend_list))

    if create_list:
      self._og.submit_bulk(create_list)
      self._logger.debug('Create %d orders', len(create_list))

    if create_list or amend_list:
      self._last_fire_time = get_timestamp()

  def manage_orders(self,
                    ask0_price,
                    bid0_price,
                    sell_theo_price,
                    buy_theo_price,
                    create_list=None,
                    amend_list=None,
                    fire=True):
    self._create_list = [] if create_list is None else create_list
    self._amend_list = [] if amend_list is None else amend_list
    self._update_orders()

    assert ask0_price > bid0_price
    ask1_price = self._tick.get_next_price(ask0_price)
    bid1_price = self._tick.get_prev_price(bid0_price)

    sell_price = None
    sell_trigger_price = None
    if sell_theo_price < bid0_price:
      assert bid0_price >= bid1_price
      sell_price = bid0_price
      sell_trigger_price = bid1_price

    buy_price = None
    buy_trigger_price = None
    if buy_theo_price >= ask0_price:
      assert ask0_price < ask1_price
      buy_price = ask0_price
      buy_trigger_price = ask1_price

    self._cancel_sell_triggered_orders(sell_theo_price, ask1_price)
    self._cancel_buy_triggered_orders(bid1_price, buy_theo_price)

    self._manage_sell_order(sell_price, sell_trigger_price)
    self._manage_buy_order(buy_price, buy_trigger_price)

    if fire:
      self.execute(self._create_list, self._amend_list)
